package utils

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base32"
	"fmt"
	"github.com/golang-jwt/jwt"
	"io"
	"strings"
	"time"
)

type ApiKeyGenerator struct {
	originalJWT string
	modelName   string
	secretKey   string
}

func NewApiKeyGenerator(userJWT, modelName, secretKey string) *ApiKeyGenerator {
	return &ApiKeyGenerator{
		originalJWT: userJWT,
		modelName:   modelName,
		secretKey:   secretKey,
	}
}

const ChatCompletionsApiKeySalt = "api-key-chat-completions"

func (g *ApiKeyGenerator) Generate() (string, error) {
	token, err := parseJWT(g.originalJWT, g.secretKey)
	if err != nil {
		return "", err
	}
	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return "", fmt.Errorf("invalid token")
	}
	newClaims := jwt.MapClaims{}
	for k, v := range claims {
		newClaims[k] = v
	}
	newClaims["model"] = g.modelName
	apiKeyJwt, err := generateAPIKeyJWT(newClaims)
	if err != nil {
		return "", fmt.Errorf("failed to generate api key jwt: %v", err)
	}
	apiKey, err := encryptJWT(apiKeyJwt, []byte(ChatCompletionsApiKeySalt))
	if err != nil {
		return "", fmt.Errorf("failed to encrypt api key jwt: %v", err)
	}
	return apiKey, nil
}
func parseJWT(originalJWT string, secretKey string) (*jwt.Token, error) {
	return jwt.Parse(originalJWT, func(token *jwt.Token) (interface{}, error) {
		return []byte(secretKey), nil
	})
}
func generateAPIKeyJWT(mc jwt.MapClaims) (string, error) {
	return jwt.NewWithClaims(jwt.SigningMethodHS256, mc).SignedString([]byte(ChatCompletionsApiKeySalt))
}

var encodingPrefix = "shy-"

func encryptJWT(jwt string, key []byte) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}
	ciphertext := gcm.Seal(nonce, nonce, []byte(jwt), nil)
	encoded := base32.StdEncoding.EncodeToString(ciphertext)
	return encodingPrefix + encoded, nil
}

func decryptJWT(shyToken string, key []byte) (string, error) {
	token := strings.TrimPrefix(shyToken, encodingPrefix)
	data, err := base32.StdEncoding.DecodeString(token)
	if err != nil {
		return "", err
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	nonceSize := gcm.NonceSize()
	if len(data) < nonceSize {
		return "", fmt.Errorf("invalid token")
	}
	nonce, ciphertext := data[:nonceSize], data[nonceSize:]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return "", err
	}
	return string(plaintext), nil
}

func CheckChatCompletionsApiKey(apiKey string) error {
	return checkApiKey(apiKey, ChatCompletionsApiKeySalt)
}

func checkApiKey(apiKey string, secretKey string) error {
	jwtStr, decryptErr := decryptJWT(apiKey, []byte(secretKey))
	if decryptErr != nil {
		return decryptErr
	}
	token, parseErr := parseJWT(jwtStr, secretKey)
	if parseErr != nil {
		return parseErr
	}
	if !token.Valid {
		return fmt.Errorf("invalid token")
	}
	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return fmt.Errorf("invalid token")
	}
	if exp, ok := claims["exp"].(float64); ok {
		if int64(exp) < time.Now().Unix() {
			return fmt.Errorf("token expired")
		}
	}

	return nil
}
